#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
by:杨锦成
vx:18873833517
time:2024.8.26
file_name:ros_initialization.py
'''

import rospy
from mavros_msgs.srv import CommandBool, SetMode, CommandLong
from mavros_msgs.msg import State, PositionTarget
from nav_msgs.msg import Odometry
from std_msgs.msg import Float32MultiArray
from geometry_msgs.msg import Twist
import tf

# 全局变量存储最新的cmd_vel消息和当前位置
cmd_vel_msg = None
current_position_x = 0.0
current_position_y = 0.0
current_position_z = 0.0  # 当前高度
current_yaw = 0.0  # 当前航向角（yaw）

# 回调函数定义
def state_callback(data, action_instance):
    action_instance.state = data  # 更新无人机状态

def odom_callback(data, action_instance):
    # 更新当前位置
    action_instance.current_position_x = data.pose.pose.position.x
    action_instance.current_position_y = data.pose.pose.position.y
    action_instance.urrent_position_z = data.pose.pose.position.z

    # 更新yaw角度
    orientation_q = data.pose.pose.orientation
    orientation_list = [orientation_q.x, orientation_q.y, orientation_q.z, orientation_q.w]
    (roll, pitch, action_instance.current_yaw) = tf.transformations.euler_from_quaternion(orientation_list)


def cmd_vel_callback(msg):
    global cmd_vel_msg
    cmd_vel_msg = msg  # 保存最新的cmd_vel消息

# ROS初始化函数
def init_ros_publishers():
    setpoint_pub = rospy.Publisher('/mavros/setpoint_raw/local', PositionTarget, queue_size=10)
    nav_goal_pub = rospy.Publisher('nav_goal', Float32MultiArray, queue_size=10)
    return setpoint_pub, nav_goal_pub

def sim_init_ros_publishers():
    setpoint_pub = rospy.Publisher('/iris_0/mavros/setpoint_raw/local', PositionTarget, queue_size=10)
    nav_goal_pub = rospy.Publisher('nav_goal', Float32MultiArray, queue_size=10)
    return setpoint_pub, nav_goal_pub

def sim_init_ros_services():
    rospy.wait_for_service('/iris_0/mavros/cmd/arming')
    rospy.wait_for_service('/iris_0/mavros/set_mode')
    rospy.wait_for_service('/iris_0/mavros/cmd/command')
    arming_service = rospy.ServiceProxy('/iris_0/mavros/cmd/arming', CommandBool)
    set_mode_service = rospy.ServiceProxy('/iris_0/mavros/set_mode', SetMode)
    command_service = rospy.ServiceProxy('/iris_0/mavros/cmd/command', CommandLong)
    return arming_service, set_mode_service, command_service

def init_ros_services():
    rospy.wait_for_service('/mavros/cmd/arming')
    rospy.wait_for_service('/mavros/set_mode')
    rospy.wait_for_service('/mavros/cmd/command')
    arming_service = rospy.ServiceProxy('/mavros/cmd/arming', CommandBool)
    set_mode_service = rospy.ServiceProxy('/mavros/set_mode', SetMode)
    command_service = rospy.ServiceProxy('/mavros/cmd/command', CommandLong)
    return arming_service, set_mode_service, command_service


def sim_init_ros_subscribers(action_instance):
    rospy.Subscriber('/iris_0/mavros/state', State, state_callback, callback_args=action_instance, queue_size=10)
    rospy.Subscriber('/iris_0/mavros/local_position/odom', Odometry, odom_callback, callback_args=action_instance, queue_size=10)
    rospy.Subscriber('cmd_vel', Twist, cmd_vel_callback, queue_size=2)

def init_ros_subscribers(action_instance):
    rospy.Subscriber('/mavros/state', State, state_callback, callback_args=action_instance, queue_size=10)
    rospy.Subscriber('/mavros/local_position/odom', Odometry, odom_callback, callback_args=action_instance, queue_size=10)
    rospy.Subscriber('cmd_vel', Twist, cmd_vel_callback, queue_size=2)

def init_ros_node(node_name):
    rospy.init_node(node_name, anonymous=True)

def get_current_yaw():
    global current_yaw
    return current_yaw

def yolo_callback(self, data):
    """
    内部回调函数，用于处理 /yolo_detect 话题的消息，更新目标X和Y坐标。
    """
    try:
        # 假设 /yolo_detect 的消息格式是 'class x: y:'
        yolo_data = data.data.split()  # 将消息按空格分隔
        self.yolo_x = float(yolo_data[1].split(":")[1])
        self.yolo_y = float(yolo_data[2].split(":")[1])
    except (IndexError, ValueError):
        rospy.logwarn("无法解析 /yolo_detect 消息，忽略此数据")


def yolo_callback(data, action_instance):
    """
    回调函数：处理来自 /AI_detect_info 话题的消息，更新目标X和Y坐标。
    
    参数:
    - data (String): 从 /AI_detect_info 话题接收到的消息，包含目标的X和Y坐标信息。
    - action_instance (Action_t): Action_t 类的实例，用于更新类的属性。
    """
    try:
        # 假设 /AI_detect_info 的消息格式是 'class: x: y; class: x: y; ...'
        if not data.data.strip():
            rospy.logwarn("收到空的 /AI_detect_info 消息")
            return

        # 按分号分割每个检测结果
        detections = data.data.split(';')
        for det in detections:
            det = det.strip()
            if not det:
                continue

            # 分割每个部分，期望格式为 'class: x: y:'
            parts = det.split()
            if len(parts) < 3:
                rospy.logwarn(f"无法解析 /AI_detect_info 消息部分: {det}")
                continue

            # 提取类别、x和y坐标
            class_part = parts[0]  # 'class:'
            x_part = parts[1]      # 'x:'
            y_part = parts[2]      # 'y:'

            # 移除冒号并获取实际值
            class_label = class_part.rstrip(':')
            x_value = x_part.rstrip(':')
            y_value = y_part.rstrip(':')

            # 检查格式
            if not class_label or not x_value or not y_value:
                rospy.logwarn(f"无法解析 /AI_detect_info 消息部分: {det}")
                continue

            # 解析x和y坐标值
            # 假设消息格式为 'QR: x:320 y:240'
            try:
                # 提取数字部分
                x = float(x_part.split(':')[1])
                y = float(y_part.split(':')[1])

                # 更新目标坐标
                action_instance.yolo_x = x
                action_instance.yolo_y = y
                rospy.logdebug(f"更新YOLO目标: class={class_label}, x={x}, y={y}")

            except (IndexError, ValueError) as e:
                rospy.logwarn(f"无法解析坐标值: {det}，错误: {e}")

    except Exception as e:
        rospy.logwarn(f"无法解析 /AI_detect_info 消息: {e}")

def odom_callback(data, action_instance):
    """
    回调函数：处理来自 /odom 话题的消息，更新无人机的当前位置和yaw角度。
    
    参数:
    - data (Odometry): 从 /odom 话题接收到的消息，包含无人机的当前位置和姿态信息。
    - action_instance (Action_t): Action_t 类的实例，用于更新类的属性。
    """
    try:
        # 更新当前位置
        action_instance.current_position_x = data.pose.pose.position.x
        action_instance.current_position_y = data.pose.pose.position.y
        action_instance.current_position_z = data.pose.pose.position.z

        # 更新yaw角度
        orientation_q = data.pose.pose.orientation
        orientation_list = [orientation_q.x, orientation_q.y, orientation_q.z, orientation_q.w]
        (roll, pitch, action_instance.current_yaw) = tf.transformations.euler_from_quaternion(orientation_list)

        rospy.logdebug(f"更新位置: x={action_instance.current_position_x}, y={action_instance.current_position_y}, z={action_instance.current_position_z}, yaw={action_instance.current_yaw}")
    except AttributeError as e:
        rospy.logerr(f"odom_callback 处理错误: {e}")